package com.vrv.snmp;

import org.snmp4j.*;
import org.snmp4j.event.ResponseEvent;
import org.snmp4j.mp.MPv1;
import org.snmp4j.mp.MPv2c;
import org.snmp4j.mp.MPv3;
import org.snmp4j.mp.SnmpConstants;
import org.snmp4j.security.SecurityModels;
import org.snmp4j.security.SecurityProtocols;
import org.snmp4j.security.USM;
import org.snmp4j.smi.*;
import org.snmp4j.transport.DefaultTcpTransportMapping;
import org.snmp4j.transport.DefaultUdpTransportMapping;
import org.snmp4j.util.MultiThreadedMessageDispatcher;
import org.snmp4j.util.ThreadPool;

import java.io.IOException;
import java.net.UnknownHostException;
import java.util.*;

public class SnmpTrapServiceImpl implements SnmpService , CommandResponder {
	private Snmp snmp = null;
	private Address targetAddress = null;
	private TransportMapping<?> transport;
	private boolean canCall = true;
	private String publicString = "public";
	private MultiThreadedMessageDispatcher dispatcher;
	private Address listenAddress;
	private ThreadPool threadPool;

	public boolean getCanCall(){
		return canCall;
	}

	public SnmpTrapServiceImpl(String ip, int port)  throws UnknownHostException, IOException{
		threadPool = ThreadPool.create("TrapPool", 2);
		dispatcher = new MultiThreadedMessageDispatcher(threadPool, new MessageDispatcherImpl());
		listenAddress = GenericAddress.parse(System.getProperty("snmp4j.listenAddress", "udp:10.139.34.137/162"));
		if (listenAddress instanceof UdpAddress) {
			transport = new DefaultUdpTransportMapping((UdpAddress) listenAddress);
		} else {
			transport = new DefaultTcpTransportMapping((TcpAddress) listenAddress);
		}
		snmp = new Snmp(dispatcher, transport);
		snmp.getMessageDispatcher().addMessageProcessingModel(new MPv1());
		snmp.getMessageDispatcher().addMessageProcessingModel(new MPv2c());
		snmp.getMessageDispatcher().addMessageProcessingModel(new MPv3());
		USM usm = new USM(SecurityProtocols.getInstance(), new OctetString(MPv3.createLocalEngineID()), 0);
		SecurityModels.getInstance().addSecurityModel(usm);
		snmp.listen();
	}

	public Object oid(String oid) throws IOException {
		CommunityTarget target = new CommunityTarget();
		target.setAddress(targetAddress);
		// 通信不成功时的重试次数
		target.setRetries(2);
		// 超时时间
		target.setTimeout(1500);
		target.setVersion(SnmpConstants.version2c);

		PDU pdu = new PDU();
		pdu.add(new VariableBinding(new OID(oid), new OctetString("SnmpTrap")));
		pdu.setType(PDU.TRAP);
		ResponseEvent respEvnt = snmp.send(pdu, target);

		if (respEvnt != null && respEvnt.getResponse() != null) {
			PDU response = respEvnt.getResponse();
			Vector<VariableBinding> recVBs = (Vector<VariableBinding>) respEvnt.getResponse().getVariableBindings();
			if(recVBs.size() == 1){
				return recVBs.elementAt(0).getVariable();
			}
		}
		return null;
	}

	public Map<String, Object> oids(String[] oids, String[] names) throws IOException{
		Map<String, Object> map = new HashMap<String, Object>();
		// 设置 target
		CommunityTarget target = new CommunityTarget();
		target.setCommunity(new OctetString(this.publicString));
		target.setAddress(targetAddress);
		// 通信不成功时的重试次数
		target.setRetries(1);
		// 超时时间
		target.setTimeout(2000);   // 有可能时而刷不出出来，timeout的时间引起来的
		target.setVersion(SnmpConstants.version2c);
		// 创建 PDU
		PDU pdu = new PDU();
		for (String oid : oids) {
			pdu.add(new VariableBinding(new OID(oid)));
		}
		pdu.setType(PDU.GET);
		// 向Agent发送PDU，并接收Response
		ResponseEvent respEvnt = snmp.send(pdu, target);
		if (respEvnt != null && respEvnt.getResponse() != null) {
			Vector<VariableBinding> recVBs = (Vector<VariableBinding>) respEvnt.getResponse().getVariableBindings();
			if (recVBs != null) {
				if (recVBs.size() > 1) {
					for (int i = 0; i < recVBs.size(); i++) {
						VariableBinding recVB = recVBs.elementAt(i);
						map.put(names[i], recVB.getVariable());
					}
				} else if (recVBs.size() == 1) {
					VariableBinding recVB = recVBs.elementAt(0);
					map.put(names[0], recVB.getVariable());
				}
			}
		}
		return map;
	}

	public List<Map<String, Object>> oidTables(String[] origOid, String[] inputOid, String[] names, int step) throws IOException{
		CommunityTarget target = new CommunityTarget();
		target.setCommunity(new OctetString(this.publicString));
		target.setAddress(targetAddress);
		// 通信不成功时的重试次数
		target.setRetries(2);
		// 超时时间
		target.setTimeout(1500);
		// 创建 PDU
		PDU pdu = new PDU();
		for (String oid : inputOid) {
			pdu.add(new VariableBinding(new OID(oid)));
		}
		target.setVersion(SnmpConstants.version2c);
		pdu.setType(PDU.GETBULK);
		pdu.setMaxRepetitions(step);
		pdu.setNonRepeaters(0);
		// 向Agent发送PDU，并接收Response
		ResponseEvent respEvnt = snmp.send(pdu, target);
		List<Map<String, Object>> result = new ArrayList<Map<String, Object>>();
		if (respEvnt != null && respEvnt.getResponse() != null) {
			Vector<VariableBinding> recVBs = (Vector<VariableBinding>) respEvnt.getResponse().getVariableBindings();
			if (recVBs != null) {
				int colSize = inputOid.length;
				int getSize = recVBs.size();
				int realStep = getSize / colSize;
				for (int i = 0; i < realStep; i++) {
					// 判断是否已越界
					VariableBinding rowIcol1 = recVBs.get(i * colSize + 0);
					if (!rowIcol1.getOid().toString().contains(origOid[0])) {
						return result;
					} else {
						Map<String, Object> rowObj = new HashMap<String, Object>();
						for (int j = 0; j < colSize; j++) {
							// i:行；j:列
							rowObj.put(names[j], recVBs.get(i * colSize + j).getVariable());
						}
						result.add(rowObj);
						if (i == realStep - 1) {
							for (int j = 0; j < colSize; j++) {
								inputOid[j] = recVBs.get(i * colSize + j).getOid().toString();
							}
						}
					}
				}
			}
		}else{
			return result;
		}
		// 递归出所有的结果
		List<Map<String, Object>> tables = oidTables(origOid, inputOid, names, step);
		result.addAll(tables);

		return result;
	}

	public void run() {
		System.out.println("----Trap Receiver run ----");
		try {
			snmp.addCommandResponder(this);
			System.out.println("--- 开始监听端口，等待Trap message  ---");
		} catch (Exception ex) {
			ex.printStackTrace();
		}
	}

	@Override
    public void close() {
        try {
            if(snmp!=null){
                snmp.close();
            }
            if(transport!=null){
                transport.close();
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

	@Override
	public List<Map<String, Object>> oidTablesAll(String[] origOid, String[] names) {
		return null;
	}

	@Override
	public List<Map<String, Object>> oidTablesWalk(String[] origOid, String[] names) {
		return null;
	}

	public static void main(String[] args) throws UnknownHostException, IOException {
		SnmpTrapServiceImpl trapReceiver = new SnmpTrapServiceImpl("192.168.118.88", 162);
		trapReceiver.run();
	}

	@Override
	public void processPdu(CommandResponderEvent event) {
		System.out.println("---- 开始解析ResponderEvent:  ----");
		if (event == null || event.getPDU() == null) {
			System.out.println("[Warn] ResponderEvent or PDU is null");
			return;
		}

		Vector<? extends VariableBinding> vbVect = event.getPDU().getVariableBindings();
		for (VariableBinding vb : vbVect) {
			System.out.println(vb.getOid() + " = " + vb.getVariable());
		}
		System.out.println("----   本次ResponderEvent 解析结束  ----");


	}
}
